//
// Created by isabella on 25-5-30.
//
#include <common.h>

template<class T>
class AVLTree {
    /**
     * 节点对象
     */
    class Node {
    public:
        int key;
        T value;
        Node *left;
        Node *right;
        int height = 1;

        Node(int k, T v) {
            key = k;
            value = v;
            left = nullptr;
            right = nullptr;
        }

        Node(int key, T value, Node *left, Node *right): key(key), value(value), left(left), right(right) {
        }
    };

    /**
     * 根节点
     */
    Node *root;

    /**
     * 求解某个节点对象的高度
     * @param n
     * @return
     */
    static int h(Node *n) {
        if (n == nullptr || n->height == 0) return 0;
        return n->height;
    }

    /**
     * 更新某个节点的高度
     * @param n
     */
    static void updateHeight(Node *n) {
        n->height = max(h(n->left), h(n->right)) + 1;
    }

    /**
     * 求解某个节点的平衡因子
     * @param n 节点对象
     * @return 平衡因子
     */
    static int bf(Node *n) {
        return h(n->left) - h(n->right);
    }

    /**
     * 右旋
     * @param node
     * @return
     */
    Node *rightRotate(Node *node) {
        if (node == nullptr) return nullptr;
        Node *a = node->left;
        Node *b = a->right;

        a->right = node;
        node->left = b;

        //有先后顺序
        updateHeight(node);
        updateHeight(a);
        return a;
    }

    /**
     * 左旋
     * @param node
     * @return
     */
    Node *leftRotate(Node *node) {
        if (node == nullptr) return nullptr;
        Node *a = node->right;
        Node *b = a->left;

        a->left = node;
        node->right = b;

        updateHeight(node);
        updateHeight(a);
        return a;
    }

    /**
     * 左右旋
     * @param node
     * @return
     */
    Node *leftRightRotate(Node *node) {
        node->left = leftRotate(node->left);
        return rightRotate(node);
    }


    /**
     * 右左旋
     * @param node
     * @return
     */
    Node *rightLeftRotate(Node *node) {
        node->right = rightRotate(node->right);
        return leftRotate(node);
    }

    /**
     * 主动平衡方法
     * @param node
     * @return
     */
    Node *balance(Node *node) {
        if (node == nullptr) return node;
        int balanceFactor = bf(node);
        if (balanceFactor > 1 && bf(node->left) >= 0) {
            //ll
            return rightRotate(node);
        } else if (balanceFactor > 1 && bf(node->left) < 0) {
            //lr
            return leftRightRotate(node);
        } else if (balanceFactor < -1 && bf(node->right) <= 0) {
            //rr
            return leftRotate(node);
        } else if (balanceFactor < -1 && bf(node->right) > 0) {
            //rl
            return rightLeftRotate(node);
        } else {
            //balanced
            return node;
        }
    }

    /**
     * 插入逻辑
     * @param node 要插入的目标节点
     * @param key 要插入的key
     * @param value 要插入的value
     * @return 插入后的原始节点
     */
    Node *doPut(Node *node, int key, T value) {
        //1.已经发现空位，插入新节点
        if (node == nullptr) return new Node(key, value);
        //2.发现当前key已经存在，更新节点
        if (key == node->key) {
            node->value = value;
            return node;
        }
        //3.未查找到正确位置，继续查找
        if (key < node->key) {
            node->left = doPut(node->left, key, value);
        } else {
            node->right = doPut(node->right, key, value);
        }
        updateHeight(node);
        return balance(node);
    }


    /**
     * 删除逻辑
     * @param node 目标节点
     * @param key 目标key
     */
    Node *doRemove(Node *node, int key) {
        //空删除现象
        if (node == nullptr) return nullptr;
        if (key < node->key) {
            node->left = doRemove(node->left, key);
        } else if (key > node->key) {
            node->right = doRemove(node->right, key);
        } else {
            //根据子节点数量分类
            if (node->left == nullptr && node->right == nullptr) {
                //无子节点
                return nullptr;
            } else if (node->left == nullptr) {
                node = node->right;
            } else if (node->right == nullptr) {
                node = node->left;
            } else {
                //有两个子节点，这里使用后继节点删除法
                Node *successor = node->right;
                while (successor->left != nullptr) {
                    successor = successor->left;
                }
                //在当前节点的右子树中，删去后继节点,再用后继节点取代当前节点
                successor->right = doRemove(node->right, successor->key);
                //整理当前节点的左子树到后继节点上，完成交换
                successor->left = node->left;
                node=successor;
            }
        }
        updateHeight(node);
        return balance(node);
    }

public:
    /**
     * 更新方法
     * @param k 待插入的key
     * @param val 待插入的value
     */
    void put(int k, T val) {
        root = doPut(root, k, val);
    }


    /**
     * 删除方法
     * @param key 待删除的key
     */
    void remove(int key) {
        root = doRemove(root, key);
    }
};

int main() {
    AVLTree<int> tree{};
    tree.put(1, 1);
    tree.put(2, 2);
    tree.put(3, 3);
    tree.put(4, 4);
    tree.put(5, 5);
    tree.remove(3);
    return 0;
}
